import socket
import struct
import json
import os
import logging
import hashlib

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

# 下载目录
download_dir = r"D:\MyCode\python_project\Python-Network-Programming\Best Practice\文件传输\client\download"
os.makedirs(download_dir, exist_ok=True)
logging.info(f"下载目录: {download_dir}")

def recv_all(sock, size):
    """可靠接收指定长度的数据"""
    data = b''
    while len(data) < size:
        remaining = size - len(data)
        chunk = sock.recv(4096 if remaining > 4096 else remaining)
        if not chunk:
            break
        data += chunk
        # 计算并显示进度
        progress = len(data) / size * 100
        print(f"\r传输进度: {progress:.2f}%", end="", flush=True)
    print()  # 换行
    return data

host = 'localhost'
port = 8080

def calculate_md5(data):
    """计算数据的MD5值"""
    md5 = hashlib.md5()
    md5.update(data)
    return md5.hexdigest()

def main():
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        s.connect((host, port))
        logging.info(f"已连接到服务器 {host}:{port}")
        
        while True:
            try:
                cmd = input('>> ').strip()
                if not cmd:
                    continue
                    
                if cmd.lower() == 'exit':
                    break
                    
                s.sendall(cmd.encode('utf-8'))
                logging.info(f"已发送命令: {cmd}")
                
                # 1. 接收报头长度
                header_size_data = recv_all(s, 4)
                if len(header_size_data) != 4:
                    logging.error("报头长度无效")
                    break
                    
                header_size = struct.unpack('>I', header_size_data)[0]
                logging.debug(f"报头大小: {header_size} 字节")
                
                # 2. 接收完整报头
                header_bytes = recv_all(s, header_size)
                if len(header_bytes) != header_size:
                    logging.error(f"报头不完整: {len(header_bytes)}/{header_size}")
                    break
                    
                try:
                    header = json.loads(header_bytes.decode('utf-8'))
                    logging.info(f"收到报头: {header}")
                except Exception as e:
                    logging.error(f"报头解析失败: {str(e)}")
                    break
                
                # 检查错误
                if 'error' in header:
                    logging.error(f"服务器错误: {header['error']}")
                    continue
                    
                # 3. 接收文件数据
                total_size = header['file_size']
                if total_size <= 0:
                    logging.warning("空文件")
                    continue
                    
                file_data = recv_all(s, total_size)
                if len(file_data) != total_size:
                    logging.error(f"数据不完整: {len(file_data)}/{total_size} 字节")
                    continue
                
                # 4. 保存文件
                filename = header['filename']
                file_path = os.path.join(download_dir, filename)
                
                with open(file_path, 'wb') as f:
                    f.write(file_data)
                
                # 验证文件大小
                saved_size = os.path.getsize(file_path)
                if saved_size != total_size:
                    logging.warning(f"文件大小不匹配: {saved_size}/{total_size}")
                else:
                    logging.info(f"文件保存成功: {file_path}")
                
                # 验证MD5 (如果有)
                if 'md5' in header and header['md5'] != '待计算':
                    file_md5 = calculate_md5(file_data)
                    if file_md5 != header['md5']:
                        logging.error(f"MD5校验失败: {file_md5} != {header['md5']}")
                    else:
                        logging.info("MD5校验通过")
                
            except KeyboardInterrupt:
                logging.info("操作中断")
                break
            except Exception as e:
                logging.error(f"处理错误: {str(e)}")
                break
    except ConnectionRefusedError:
        logging.error("无法连接到服务器")
    except Exception as e:
        logging.error(f"连接错误: {str(e)}")
    finally:
        s.close()
        logging.info("连接已关闭")

if __name__ == '__main__':
    main()